{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "76fba87d",
   "metadata": {},
   "source": [
    "| [02_lexical_analysis/02_Bert实体识别.ipynb](https://github.com/shibing624/nlp-tutorial/blob/main/07_information_extraction/02_Bert实体识别.ipynb)  | Bert实体识别  |[Open In Colab](https://colab.research.google.com/github/shibing624/nlp-tutorial/blob/main/07_information_extraction/02_Bert实体识别.ipynb) |\n",
    "\n",
    "# Bert实体识别\n",
    "适用于品牌、人名、地址名称识别，序列标注任务解决方法。\n",
    "\n",
    "基于transformers的预训练模型，识别人名实体，模型地址：https://huggingface.co/models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7da5eba5",
   "metadata": {
    "pycharm": {
     "is_executing": true,
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple\n",
      "Requirement already satisfied: transformers in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (4.11.3)\n",
      "Requirement already satisfied: seqeval in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (1.2.2)\n",
      "Requirement already satisfied: scikit-learn>=0.21.3 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from seqeval) (1.0.2)\n",
      "Requirement already satisfied: numpy>=1.14.0 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from seqeval) (1.20.1)\n",
      "Requirement already satisfied: scipy>=1.1.0 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from scikit-learn>=0.21.3->seqeval) (1.6.2)\n",
      "Requirement already satisfied: joblib>=0.11 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from scikit-learn>=0.21.3->seqeval) (1.0.1)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from scikit-learn>=0.21.3->seqeval) (2.1.0)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from transformers) (2021.11.2)\n",
      "Requirement already satisfied: huggingface-hub>=0.0.17 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from transformers) (0.0.19)\n",
      "Requirement already satisfied: filelock in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from transformers) (3.4.2)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from transformers) (5.4.1)\n",
      "Requirement already satisfied: sacremoses in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from transformers) (0.0.45)\n",
      "Requirement already satisfied: tokenizers<0.11,>=0.10.1 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from transformers) (0.10.3)\n",
      "Requirement already satisfied: packaging>=20.0 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from transformers) (20.9)\n",
      "Requirement already satisfied: requests in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from transformers) (2.25.1)\n",
      "Requirement already satisfied: tqdm>=4.27 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from transformers) (4.62.3)\n",
      "Requirement already satisfied: typing-extensions in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from huggingface-hub>=0.0.17->transformers) (3.7.4.3)\n",
      "Requirement already satisfied: pyparsing>=2.0.2 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from packaging>=20.0->transformers) (2.4.7)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from requests->transformers) (2021.10.8)\n",
      "Requirement already satisfied: chardet<5,>=3.0.2 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from requests->transformers) (4.0.0)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from requests->transformers) (1.26.4)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from requests->transformers) (2.10)\n",
      "Requirement already satisfied: click in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from sacremoses->transformers) (7.1.2)\n",
      "Requirement already satisfied: six in /Users/xuming/opt/anaconda3/lib/python3.8/site-packages (from sacremoses->transformers) (1.15.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install transformers seqeval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6c650d4d",
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'entity': 'LABEL_8', 'score': 0.9998587, 'index': 1, 'word': '王', 'start': 0, 'end': 1}, {'entity': 'LABEL_5', 'score': 0.99971896, 'index': 2, 'word': '宏', 'start': 1, 'end': 2}, {'entity': 'LABEL_5', 'score': 0.99990284, 'index': 3, 'word': '伟', 'start': 2, 'end': 3}, {'entity': 'LABEL_2', 'score': 0.9999732, 'index': 4, 'word': '来', 'start': 3, 'end': 4}, {'entity': 'LABEL_2', 'score': 0.9999744, 'index': 5, 'word': '自', 'start': 4, 'end': 5}, {'entity': 'LABEL_1', 'score': 0.99971133, 'index': 6, 'word': '北', 'start': 5, 'end': 6}, {'entity': 'LABEL_4', 'score': 0.99962795, 'index': 7, 'word': '京', 'start': 6, 'end': 7}, {'entity': 'LABEL_2', 'score': 0.9999791, 'index': 8, 'word': '，', 'start': 7, 'end': 8}, {'entity': 'LABEL_2', 'score': 0.9999846, 'index': 9, 'word': '是', 'start': 8, 'end': 9}, {'entity': 'LABEL_2', 'score': 0.9999859, 'index': 10, 'word': '个', 'start': 9, 'end': 10}, {'entity': 'LABEL_2', 'score': 0.9999852, 'index': 11, 'word': '警', 'start': 10, 'end': 11}, {'entity': 'LABEL_2', 'score': 0.99998474, 'index': 12, 'word': '察', 'start': 11, 'end': 12}, {'entity': 'LABEL_2', 'score': 0.9999855, 'index': 13, 'word': '，', 'start': 12, 'end': 13}, {'entity': 'LABEL_2', 'score': 0.9999858, 'index': 14, 'word': '喜', 'start': 13, 'end': 14}, {'entity': 'LABEL_2', 'score': 0.9999869, 'index': 15, 'word': '欢', 'start': 14, 'end': 15}, {'entity': 'LABEL_2', 'score': 0.999984, 'index': 16, 'word': '去', 'start': 15, 'end': 16}, {'entity': 'LABEL_1', 'score': 0.997468, 'index': 17, 'word': '王', 'start': 16, 'end': 17}, {'entity': 'LABEL_4', 'score': 0.99714166, 'index': 18, 'word': '府', 'start': 17, 'end': 18}, {'entity': 'LABEL_4', 'score': 0.9959686, 'index': 19, 'word': '井', 'start': 18, 'end': 19}, {'entity': 'LABEL_2', 'score': 0.99992996, 'index': 20, 'word': '游', 'start': 19, 'end': 20}, {'entity': 'LABEL_2', 'score': 0.9999782, 'index': 21, 'word': '玩', 'start': 20, 'end': 21}, {'entity': 'LABEL_2', 'score': 0.999975, 'index': 22, 'word': '儿', 'start': 21, 'end': 22}, {'entity': 'LABEL_2', 'score': 0.9999846, 'index': 23, 'word': '。', 'start': 22, 'end': 23}]\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "import torch\n",
    "from transformers import AutoModelForTokenClassification, AutoTokenizer\n",
    "from transformers import pipeline\n",
    "\n",
    "os.environ[\"KMP_DUPLICATE_LIB_OK\"] = \"TRUE\"\n",
    "model_name = \"shibing624/bert4ner-base-chinese\"\n",
    "\n",
    "nlp = pipeline(\"ner\", model=model_name, tokenizer=model_name)\n",
    "sequence = \"王宏伟来自北京，是个警察，喜欢去王府井游玩儿。\"\n",
    "print(nlp(sequence))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04fff3bd",
   "metadata": {},
   "source": [
    "不使用pipline，自己编写NER任务预测代码："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fffd3640",
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "王宏伟来自北京，是个警察，喜欢去王府井游玩儿。\n",
      "[('王', 'B-PER'), ('宏', 'I-PER'), ('伟', 'I-PER'), ('来', 'O'), ('自', 'O'), ('北', 'B-LOC'), ('京', 'I-LOC'), ('，', 'O'), ('是', 'O'), ('个', 'O'), ('警', 'O'), ('察', 'O'), ('，', 'O'), ('喜', 'O'), ('欢', 'O'), ('去', 'O'), ('王', 'B-LOC'), ('府', 'I-LOC'), ('井', 'I-LOC'), ('游', 'O'), ('玩', 'O'), ('儿', 'O'), ('。', 'O')]\n",
      "Sentence entity:\n",
      "[('王宏伟', 'PER'), ('北京', 'LOC'), ('王府井', 'LOC')]\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForTokenClassification\n",
    "from seqeval.metrics.sequence_labeling import get_entities\n",
    "\n",
    "os.environ[\"KMP_DUPLICATE_LIB_OK\"] = \"TRUE\"\n",
    "\n",
    "# Load model from HuggingFace Hub\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"shibing624/bert4ner-base-chinese\")\n",
    "model = AutoModelForTokenClassification.from_pretrained(\"shibing624/bert4ner-base-chinese\")\n",
    "label_list = ['I-ORG', 'B-LOC', 'O', 'B-ORG', 'I-LOC', 'I-PER', 'B-TIME', 'I-TIME', 'B-PER']\n",
    "\n",
    "sentence = \"王宏伟来自北京，是个警察，喜欢去王府井游玩儿。\"\n",
    "\n",
    "\n",
    "def get_entity(sentence):\n",
    "    tokens = tokenizer.tokenize(sentence)\n",
    "    inputs = tokenizer.encode(sentence, return_tensors=\"pt\")\n",
    "    with torch.no_grad():\n",
    "        outputs = model(inputs).logits\n",
    "    predictions = torch.argmax(outputs, dim=2)\n",
    "    char_tags = [(token, label_list[prediction]) for token, prediction in zip(tokens, predictions[0].numpy())][1:-1]\n",
    "    print(sentence)\n",
    "    print(char_tags)\n",
    "\n",
    "    pred_labels = [i[1] for i in char_tags]\n",
    "    entities = []\n",
    "    line_entities = get_entities(pred_labels)\n",
    "    for i in line_entities:\n",
    "        word = sentence[i[1]: i[2] + 1]\n",
    "        entity_type = i[0]\n",
    "        entities.append((word, entity_type))\n",
    "\n",
    "    print(\"Sentence entity:\")\n",
    "    print(entities)\n",
    "\n",
    "\n",
    "get_entity(sentence)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b74bc373",
   "metadata": {},
   "source": [
    "本节完。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "700b6906",
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
